from concurrent.futures import ThreadPoolExecutor, as_completed

from Crypto.Cipher import AES

import requests
import time
import os
import sys
import shutil


def decrypt(src, key):
    aes = AES.new(key, AES.MODE_CBC, key)
    des = aes.decrypt(src)
    return des

def dld(url):
    for i in range(3):
        ret = requests.get(url, timeout=600)
        if ret.ok:
            return ret.content
        time.sleep(1)

def write(url, key, file_name):
    start = time.time()
    data = dld(url)
    data = decrypt(data, key)
    with open(file_name, "wb") as fp:
        fp.write(data)
    print(url, time.time()-start)

executor = ThreadPoolExecutor(max_workers=100)


def start_dld(m3u8, tmp_path):
    with open(m3u8, "r") as fp:
        m3u8 = fp.readlines()
    all_task = []
    file_list = []
    i = 0
    for line in m3u8:
        line = line.strip()
        if "METHOD" in line:
            url = line.split("URI=")[-1].strip("\"")
            key = dld(url)
            print(key) 
        if line.startswith("http"):
            i += 1
            tmp_name = os.path.join(tmp_path, str(i)) 
            t = executor.submit(write, line, key, tmp_name)
            all_task.append(t)
            file_list.append(tmp_name)
    for future in as_completed(all_task):
        future.result()
    return file_list


def hebing(m3u8, file_name):
    print(m3u8, file_name)
    try:
        tmp_path = os.path.join(os.path.dirname(file_name), "." + os.path.basename(file_name))
        try:
            os.makedirs(tmp_path)
        except:
            pass

        file_list = start_dld(m3u8, tmp_path)

        with open(file_name, "wb") as mp4_fp:
            for _file in file_list:
                mp4_fp.write(open(_file, "rb").read())
    finally:
        shutil.rmtree(tmp_path)

    return True

if __name__ == "__main__":
    hebing(r"\\192.168.123.1\sda1\3\index.m3u8.mp4", r"d:\mp4\1.mp4")